/*
 * Copyright (c) Kumo Inc. and affiliates.
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <melon/json/bser/bser.h>

#include <melon/io/cursor.h>

using namespace melon;
using melon::bser::serialization_opts;
using melon::io::QueueAppender;

namespace melon {
    namespace bser {
        const uint8_t kMagic[2] = {0, 1};

        static void bserEncode(
            Dynamic const &dyn,
            QueueAppender &appender,
            const serialization_opts &opts);

        serialization_opts::serialization_opts()
            : sort_keys(false), growth_increment(8192) {
        }

        static const Dynamic *getTemplate(
            const serialization_opts &opts, Dynamic const &dynArray) {
            if (!opts.templates.has_value()) {
                return nullptr;
            }
            const auto &templates = opts.templates.value();
            const auto it = templates.find(&dynArray);
            if (it == templates.end()) {
                return nullptr;
            }
            return &it->second;
        }

        static void bserEncodeInt(int64_t ival, QueueAppender &appender) {
            /* Return the smallest size int that can store the value */
            auto size =
            ((ival == ((int8_t) ival))
                 ? 1
                 : (ival == ((int16_t) ival))
                       ? 2
                       : (ival == ((int32_t) ival))
                             ? 4
                             : 8);

            switch (size) {
                case 1:
                    appender.write((int8_t) BserType::Int8);
                    appender.write(int8_t(ival));
                    return;
                case 2:
                    appender.write((int8_t) BserType::Int16);
                    appender.write(int16_t(ival));
                    return;
                case 4:
                    appender.write((int8_t) BserType::Int32);
                    appender.write(int32_t(ival));
                    return;
                case 8:
                    appender.write((int8_t) BserType::Int64);
                    appender.write(ival);
                    return;
                default:
                    throw std::runtime_error("impossible integer size");
            }
        }

        static void bserEncodeString(melon::StringPiece str, QueueAppender &appender) {
            appender.write((int8_t) BserType::String);
            bserEncodeInt(int64_t(str.size()), appender);
            appender.push((uint8_t *) str.data(), str.size());
        }

        static void bserEncodeArraySimple(
            Dynamic const &dyn,
            QueueAppender &appender,
            const serialization_opts &opts) {
            appender.write((int8_t) BserType::Array);
            bserEncodeInt(int64_t(dyn.size()), appender);
            for (const auto &ele: dyn) {
                bserEncode(ele, appender, opts);
            }
        }

        static void bserEncodeArray(
            Dynamic const &dyn,
            QueueAppender &appender,
            const serialization_opts &opts) {
            auto templ = getTemplate(opts, dyn);
            if (MELON_UNLIKELY(templ != nullptr)) {
                appender.write((int8_t) BserType::Template);

                // Emit the list of property names
                bserEncodeArraySimple(*templ, appender, opts);

                // The number of objects in the array
                bserEncodeInt(int64_t(dyn.size()), appender);

                // For each object in the array
                for (const auto &ele: dyn) {
                    // For each key in the template
                    for (const auto &name: *templ) {
                        if (auto found = ele.get_ptr(name)) {
                            if (found->isNull()) {
                                // Prefer to Skip rather than encode a null value for
                                // compatibility with the other bser implementations
                                appender.write((int8_t) BserType::Skip);
                            } else {
                                bserEncode(*found, appender, opts);
                            }
                        } else {
                            appender.write((int8_t) BserType::Skip);
                        }
                    }
                }
                return;
            }

            bserEncodeArraySimple(dyn, appender, opts);
        }

        static void bserEncodeObject(
            Dynamic const &dyn,
            QueueAppender &appender,
            const serialization_opts &opts) {
            appender.write((int8_t) BserType::Object);
            bserEncodeInt(int64_t(dyn.size()), appender);

            if (opts.sort_keys) {
                std::vector<std::pair<Dynamic, Dynamic> > sorted(
                    dyn.items().begin(), dyn.items().end());
                std::sort(sorted.begin(), sorted.end());
                for (const auto &item: sorted) {
                    bserEncode(item.first, appender, opts);
                    bserEncode(item.second, appender, opts);
                }
            } else {
                for (const auto &item: dyn.items()) {
                    bserEncode(item.first, appender, opts);
                    bserEncode(item.second, appender, opts);
                }
            }
        }

        static void bserEncode(
            Dynamic const &dyn,
            QueueAppender &appender,
            const serialization_opts &opts) {
            switch (dyn.type()) {
                case Dynamic::Type::NULLT:
                    appender.write((int8_t) BserType::Null);
                    return;
                case Dynamic::Type::BOOL:
                    appender.write(
                        (int8_t) (dyn.getBool() ? BserType::True : BserType::False));
                    return;
                case Dynamic::Type::DOUBLE: {
                    double dval = dyn.getDouble();
                    appender.write((int8_t) BserType::Real);
                    appender.write(dval);
                    return;
                }
                case Dynamic::Type::INT64:
                    bserEncodeInt(dyn.getInt(), appender);
                    return;
                case Dynamic::Type::OBJECT:
                    bserEncodeObject(dyn, appender, opts);
                    return;
                case Dynamic::Type::ARRAY:
                    bserEncodeArray(dyn, appender, opts);
                    return;
                case Dynamic::Type::STRING:
                    bserEncodeString(dyn.getString(), appender);
                    return;
            }
        }

        std::unique_ptr<melon::IOBuf> toBserIOBuf(
            melon::Dynamic const &dyn, const serialization_opts &opts) {
            IOBufQueue q(IOBufQueue::cacheChainLength());
            uint8_t hdrbuf[sizeof(kMagic) + 1 + sizeof(int64_t)];

            // Reserve some headroom for the overall PDU size; we'll fill this in
            // after we've serialized the data and know the length
            auto firstbuf = IOBuf::create(opts.growth_increment);
            firstbuf->advance(sizeof(hdrbuf));
            q.append(std::move(firstbuf));

            // encode the value
            QueueAppender appender(&q, opts.growth_increment);
            bserEncode(dyn, appender, opts);

            // compute the length
            auto len = q.chainLength();
            if (len > uint64_t(std::numeric_limits<int64_t>::max())) {
                throw std::range_error(melon::to<std::string>(
                    "serialized data size ", len, " is too large to represent as BSER"));
            }

            // This is a bit verbose, but it computes a header that is appropriate
            // to the size of the serialized data

            memcpy(hdrbuf, kMagic, sizeof(kMagic));
            size_t hdrlen = sizeof(kMagic) + 1;
            auto magicptr = hdrbuf + sizeof(kMagic);
            auto lenptr = hdrbuf + hdrlen;

            if (len > uint64_t(std::numeric_limits<int32_t>::max())) {
                *magicptr = (int8_t) BserType::Int64;
                *(int64_t *) lenptr = (int64_t) len;
                hdrlen += sizeof(int64_t);
            } else if (len > uint64_t(std::numeric_limits<int16_t>::max())) {
                *magicptr = (int8_t) BserType::Int32;
                *(int32_t *) lenptr = (int32_t) len;
                hdrlen += sizeof(int32_t);
            } else if (len > uint64_t(std::numeric_limits<int8_t>::max())) {
                *magicptr = (int8_t) BserType::Int16;
                *(int16_t *) lenptr = (int16_t) len;
                hdrlen += sizeof(int16_t);
            } else {
                *magicptr = (int8_t) BserType::Int8;
                *(int8_t *) lenptr = (int8_t) len;
                hdrlen += sizeof(int8_t);
            }

            // and place the data in the headroom
            q.prepend(hdrbuf, hdrlen);

            return q.move();
        }

        kmstring toBser(Dynamic const &dyn, const serialization_opts &opts) {
            auto buf = toBserIOBuf(dyn, opts);
            return buf->moveToKmString();
        }
    } // namespace bser
} // namespace melon

/* vim:ts=2:sw=2:et:
 */
